import random
import numpy as np

class Game:
    def __init__(self,group,agent1,agent2):
        self.agent1=agent1
        self.agent2=agent2
        self.group=group
        #self.winner=None
    def run(self):
        move1=0
        move2=0
        agent_1_move={}
        agent_2_move={}
        for i in self.group:
            agent_1_move[i]=0
            agent_2_move[i]=0
        for i in range(100):
            move=self.agent1.chooseMove()
            agent_1_move[move]+=1
            move1+=move
            move=self.agent2.chooseMove()
            agent_2_move[move]+=1
            move2+=move
        if move1<move2:
            self.winner='agent2'
            return agent_2_move,agent_1_move #return win, lose
        elif move1>move2:
            self.winner='agent1'
            return agent_1_move,agent_2_move
        else:
            self.winner='agent1'
            return agent_1_move,agent_2_move #平局也返回agent1赢，不影响结果
    def updatePolicy(self,policy):
        self.agent1.policy=policy



class agent:
    def __init__(self,group,policy):
        self.policy=policy
        self.group=group
    def chooseMove(self):
        return random.choices(self.group,weights=self.policy)[0]


def normlize(policy):
    policy[policy<0]=0
    return policy/np.sum(policy)

def main():
    #游戏随机从group中选一个数字，然后双方选100次后比总和的大小
    group=[1,2,3,4,5]
    policy=[.2,.2,.2,.2,.2] #初始策略平均获取
    policy1=policy
    policy2=policy
    agent1=agent(group,policy1)
    agent2=agent(group,policy2)
    game=Game(group,agent1,agent2)
    round=1000 #玩round局
    for i in range(round):
        win,lose=game.run()
        net=np.array(list(win.values()))-np.array(list(lose.values()))
        lr=0.001
        policy1_=np.array(policy1)+lr*net
        policy1=normlize(policy1_).tolist()
        game.updatePolicy(policy1)
    print(policy1)

if __name__ == "__main__":
    main()
